Skip to content

Avoid dynamic dispatch in back-end IR code - #1071

Merged
maleadt merged 6 commits into
mainfrom
tb/typed-ir-calls
Oct 30, 2026
Merged

maleadt merged 6 commits into
mainfrom
tb/typed-ir-calls

Conversation

@maleadt

@maleadt maleadt commented Oct 28, 2026

Copy link
Copy Markdown
Member

Compiling the first kernel of a session JIT-compiles a lot of LLVM.jl, GPUCompiler and Base host code, much of it because GPUCompiler's back-end code calls it through dynamic dispatch: values come out of untyped containers (classify_arguments' Vector{Any}, Metal's untyped kernel_intrinsics dict, comprehensions whose element type widens at run time), IR objects get splatted into vector literals, and the atomics helpers are specialized on every instruction kind they see. Many of the callees are only inferred in an upstream package image, without native code, and a back-end's own image can't add native code to those (JuliaLang/julia#58465), so this happens in every session.

These commits make those call sites type-stable, or stop the specialization where it doesn't buy anything: a typed vector from classify_arguments, vectors of IR objects built with collect/pushfirst!/concatenation instead of splatting, typed atomics worklists for PTX, Metal and SPIR-V, @nospecialize on the local-atomics and shared atomics helpers, and a few type assertions in Metal's argument metadata. Behavior is unchanged.

Together with #1070, on Julia 1.13 the first new kernel takes 939 instead of 1470 ms with CUDA.jl on an RTX 5080 (JIT-compiled methods 156 → 124), and 233 instead of 270 ms with Metal.jl on an M1 (56 → 47).

What remains is mostly in LLVM.jl (accessors specialized on the concrete instruction type, first_use dispatched inside iterate(::ValueUseSet)) and in the back-ends' own packages.

The arguments were collected in a `Vector{Any}`, so every use of their
fields in the back-ends (Metal's argument metadata and parameter
attributes, `irgen`'s by-value rewriting) dispatched dynamically, and
so did the LLVM.jl and Base methods called with them. Many of those
are only inferred in an upstream package image, without native code,
and a back-end's own image can't add native code to them
(JuliaLang/julia#58465), so they were compiled in every session.
Splatting a collection of LLVM values or types into a vector literal
dispatches on the run-time types of the elements, compiling a method
for every combination of instruction and type kinds encountered. Build
these vectors with `collect`, `pushfirst!` or by concatenating
concretely typed collections instead. Also assert the types of the
kernel state getter call and of the caller instruction `backtrace`
moves to, which kept every property access in its loop dynamic.
The comprehension collecting the atomics widened its element type at
run time, so it was inferred as an untyped vector and every helper
called on its elements was dispatched dynamically. Collect into a
`Vector{LLVM.Instruction}`, don't specialize the helpers on the kind of
instruction or pointer (which compiled each of them again for every
kind), and give the other vector literals built from IR objects an
element type.
As for PTX, collect the atomics to lower into a typed vector, don't
specialize the helpers on the kind of instruction or pointer, and build
the argument and parameter-type vectors of the AIR atomic intrinsics by
appending to typed literals rather than by splatting into untyped ones.
Give the table of kernel intrinsics that become arguments concrete key
and value types, which made all of `add_input_arguments!` dynamic. When
emitting argument metadata, assert the types of sizes and address
spaces computed from untyped values, so that the `Int32` and
`ConstantInt` calls on them are resolved statically.
Like for PTX and Metal, the untyped comprehension made every call on
the collected instructions dispatch dynamically.
Their callers only know they have an `LLVM.Instruction` or
`LLVM.Value`, so every call dispatched dynamically and compiled the
helper again for each concrete kind of instruction or pointer it met.
With `@nospecialize`, a single method instance serves all of them and
the calls resolve statically.
@codecov

codecov Bot commented Oct 29, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 96.22642% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.46%. Comparing base (daab6f8) to head (9085f99).

Files with missing lines Patch % Lines
src/irgen.jl 71.42% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1071      +/-   ##
==========================================
- Coverage   92.49%   92.46%   -0.03%     
==========================================
  Files          31       31              
  Lines        7313     7315       +2     
==========================================
  Hits         6764     6764              
- Misses        549      551       +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@maleadt
maleadt merged commit b90c3cb into main Oct 30, 2026
63 of 64 checks passed
@maleadt
maleadt deleted the tb/typed-ir-calls branch October 30, 2026 09:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant